import os
import numpy as np
import json


class DataReader():
    def __init__(self,data_dir,ids_path,feature_len,seq_max_len):
        self._data_dir = data_dir   # 数据集目录
        self._ids_path=ids_path      # 标签路径
        self._feature_len=feature_len   # 特征长度
        self._seq_max_len=seq_max_len   # 序列最大长度

        self.label2id=None
        self.id2label=None
        self.data=None
        self.label=None
        self.data_length=None

        self.loadLabelId()
        self.loadData()

    def loadLabelId(self):
        self.label2id,self.id2label=getJsonData(self._ids_path,
                                                ['label2id','id2label'])

    def loadData(self):
        # 加载数据
        word_list=list(self.label2id.keys())
        self.data,self.label=getAllData(self._data_dir,
                                        self.label2id,
                                        ["lip_features"],
                                        word_list)
        self.label=convert_to_one_hot(self.label)
        self.data=paddingData(self.data,self._feature_len,self._seq_max_len)
        self.data_length=len(self.label)     # 数据集长度

    def getBatchData(self,batch_size):
        # 获取批量数据
        cursor = 0
        flag=True
        while flag:
            batch_data=[]
            batch_label=[]
            for _ in range(batch_size):
                batch_data.append(self.data[cursor])
                batch_label.append(self.label[cursor])
                cursor+=1
                if cursor>=self.data_length:
                    cursor=0
                    flag=False
                    break
            yield batch_data,batch_label

    def getAllData(self,):
        # 返回所有的数据
        return self.data,self.label


def convert_to_one_hot(data):
    # 转化独热向量
    data=np.array(data)
    C=data.max()+1    # 类别
    return np.eye(C)[data.reshape(-1)]


def paddingData(features,feature_len,sequence_len):
    # 对数据进行填充
    zero=[[0]*feature_len]
    for feature in features:
        feature.extend((sequence_len-len(feature))*zero)
    return features


def getJsonData(json_path,feature_name):
    # 获取特征数据
    with open(json_path,"r",encoding="utf-8") as f:
        d=json.load(f)
        features=[]
        for name in feature_name:
            features.append(d[name])
    return features


def getAllData(json_dir,label2id,feature_name,extract_word_list):
    # 获取所有特征数据
    files=os.listdir(json_dir)  # 获取目录下所有文件名
    features_tmp=[]
    for file in files:
        word_label=file.rstrip(".json")      # 获取标签
        if word_label in extract_word_list or extract_word_list is None:
            feature=getJsonData(os.path.join(json_dir,file),feature_name)
            feature=[(i,label2id[word_label])for i in feature[0]]
            features_tmp.extend(feature)
    np.random.shuffle(features_tmp)     # 打乱数据集
    features=[i[0] for i in features_tmp]
    labels=[i[1] for i in features_tmp]
    return features,labels


if __name__=="__main__":
    from LSTMModel.config import lstm_config
    data_reader=DataReader(lstm_config.train_dir,
                           lstm_config.ids_path,0,0)
    data_reader.loadLabelId()
    data_reader.loadData()
    num_count=[]
    for video in data_reader.data:
            num_count.append(len(video))

    # 输出信息
    print("训练文件总数量：{}".format(data_reader.data_length))
    print("最大输入维度：{}".format(max(num_count)))
    print("平均输入维度：{}".format(sum(num_count) / data_reader.data_length))

    # 可视化
    # matplotlib.rcParams['font.sans-serif'] = ['SimHei']
    # matplotlib.rcParams['font.family'] = 'sans-serif'
    # plt.hist(num_count,30,facecolor='g')
    # plt.xlabel("file length")
    # plt.ylabel("frequency")
    # plt.axis([0,30,0,2000])
    # plt.show()


